"""
幂等性管理器 - 确保任务执行的幂等性，防止重复执行和竞态条件
"""
import hashlib
import time
from typing import Dict, Any, Optional, Set
from datetime import datetime, timedelta
from dataclasses import dataclass
from enum import Enum

from utils.logger import Logger
from middleware.redis_adapter import get_redis_adapter
from utils.exceptions import IdempotencyError

logger = get_logger("idempotency_manager")


class IdempotencyState(Enum):
    """幂等性状态枚举"""
    PENDING = "pending"      # 处理中
    COMPLETED = "completed"  # 已完成
    FAILED = "failed"       # 已失败
    EXPIRED = "expired"     # 已过期


@dataclass
class IdempotencyRecord:
    """幂等性记录"""
    request_id: str
    state: IdempotencyState
    result: Optional[Any] = None
    error_message: Optional[str] = None
    created_at: datetime = None
    updated_at: datetime = None
    attempt_count: int = 0
    max_attempts: int = 3
    ttl: int = 3600  # 默认1小时过期


class IdempotencyManager:
    """幂等性管理器"""
    
    def __init__(self, redis_client=None):
        self.redis_client = redis_client or get_redis_adapter()
        self.key_prefix = "idempotency:"
        self.lock_prefix = "idempotency_lock:"
        self.default_ttl = 3600  # 1小时
        self.max_attempts = 3
        
    def _generate_request_id(self, **kwargs) -> str:
        """基于请求参数生成request_id"""
        # 将参数排序并序列化
        sorted_params = sorted(kwargs.items())
        param_str = ",".join([f"{k}={v}" for k, v in sorted_params])
        
        # 使用SHA256生成哈希
        return hashlib.sha256(param_str.encode()).hexdigest()[:16]
    
    def _get_record_key(self, request_id: str) -> str:
        """获取记录键"""
        return f"{self.key_prefix}{request_id}"
    
    def _get_lock_key(self, request_id: str) -> str:
        """获取锁键"""
        return f"{self.lock_prefix}{request_id}"
    
    async def check_idempotency(self, request_id: str) -> Optional[IdempotencyRecord]:
        """检查幂等性状态"""
        try:
            key = self._get_record_key(request_id)
            record_data = self.redis_client.get(key)
            
            if not record_data:
                return None
            
            import json
            record_dict = json.loads(record_data.decode('utf-8'))
            record = IdempotencyRecord(**record_dict)
            
            # 检查是否过期
            if datetime.fromisoformat(record.updated_at) < datetime.utcnow() - timedelta(seconds=record.ttl):
                record.state = IdempotencyState.EXPIRED
                
            return record
            
        except Exception as e:
            logger.error(f"检查幂等性失败: {request_id}, error={str(e)}")
            return None
    
    async def acquire_processing_lock(self, request_id: str, timeout: int = 30) -> bool:
        """获取处理锁（防止并发执行）"""
        try:
            lock_key = self._get_lock_key(request_id)
            
            # 使用Redis的SET NX EX命令实现分布式锁
            result = self.redis_client.set(
                lock_key, 
                "locked", 
                nx=True, 
                ex=timeout
            )
            
            if result:
                logger.info(f"获取处理锁成功: {request_id}")
                return True
            else:
                logger.warning(f"获取处理锁失败（可能已被其他进程获取）: {request_id}")
                return False
                
        except Exception as e:
            logger.error(f"获取处理锁异常: {request_id}, error={str(e)}")
            return False
    
    async def release_processing_lock(self, request_id: str) -> bool:
        """释放处理锁"""
        try:
            lock_key = self._get_lock_key(request_id)
            self.redis_client.delete(lock_key)
            logger.info(f"释放处理锁成功: {request_id}")
            return True
            
        except Exception as e:
            logger.error(f"释放处理锁异常: {request_id}, error={str(e)}")
            return False
    
    async def start_processing(self, request_id: str, **kwargs) -> bool:
        """开始处理（创建幂等性记录）"""
        try:
            # 检查是否已存在记录
            existing_record = await self.check_idempotency(request_id)
            if existing_record:
                if existing_record.state == IdempotencyState.PENDING:
                    raise IdempotencyError(f"请求正在处理中: {request_id}")
                elif existing_record.state == IdempotencyState.COMPLETED:
                    logger.info(f"请求已处理完成，返回缓存结果: {request_id}")
                    return False  # 表示不需要重新处理
                elif existing_record.state == IdempotencyState.FAILED:
                    if existing_record.attempt_count >= existing_record.max_attempts:
                        raise IdempotencyError(f"请求已失败且达到最大重试次数: {request_id}")
            
            # 获取处理锁
            if not await self.acquire_processing_lock(request_id):
                raise IdempotencyError(f"无法获取处理锁，可能其他进程正在处理: {request_id}")
            
            # 创建新的处理记录
            record = IdempotencyRecord(
                request_id=request_id,
                state=IdempotencyState.PENDING,
                attempt_count=existing_record.attempt_count + 1 if existing_record else 1,
                max_attempts=kwargs.get('max_attempts', self.max_attempts),
                ttl=kwargs.get('ttl', self.default_ttl),
                created_at=datetime.utcnow(),
                updated_at=datetime.utcnow()
            )
            
            # 保存记录
            await self._save_record(record)
            logger.info(f"开始处理请求: {request_id}, attempt={record.attempt_count}")
            return True
            
        except IdempotencyError:
            raise
        except Exception as e:
            logger.error(f"开始处理异常: {request_id}, error={str(e)}")
            raise IdempotencyError(f"开始处理失败: {str(e)}")
    
    async def complete_processing(self, request_id: str, result: Any) -> bool:
        """完成处理"""
        try:
            record = await self.check_idempotency(request_id)
            if not record:
                logger.warning(f"完成处理时发现记录不存在: {request_id}")
                return False
            
            record.state = IdempotencyState.COMPLETED
            record.result = result
            record.updated_at = datetime.utcnow()
            
            await self._save_record(record)
            await self.release_processing_lock(request_id)
            
            logger.info(f"处理完成: {request_id}")
            return True
            
        except Exception as e:
            logger.error(f"完成处理异常: {request_id}, error={str(e)}")
            return False
    
    async def fail_processing(self, request_id: str, error_message: str) -> bool:
        """处理失败"""
        try:
            record = await self.check_idempotency(request_id)
            if not record:
                logger.warning(f"处理失败时发现记录不存在: {request_id}")
                return False
            
            record.state = IdempotencyState.FAILED
            record.error_message = error_message
            record.updated_at = datetime.utcnow()
            
            await self._save_record(record)
            await self.release_processing_lock(request_id)
            
            logger.error(f"处理失败: {request_id}, error={error_message}")
            return True
            
        except Exception as e:
            logger.error(f"处理失败异常: {request_id}, error={str(e)}")
            return False
    
    async def _save_record(self, record: IdempotencyRecord) -> bool:
        """保存记录"""
        try:
            key = self._get_record_key(record.request_id)
            import json
            
            # 将dataclass转换为dict
            record_dict = {
                'request_id': record.request_id,
                'state': record.state.value,
                'result': record.result,
                'error_message': record.error_message,
                'created_at': record.created_at.isoformat() if record.created_at else None,
                'updated_at': record.updated_at.isoformat() if record.updated_at else None,
                'attempt_count': record.attempt_count,
                'max_attempts': record.max_attempts,
                'ttl': record.ttl
            }
            
            record_json = json.dumps(record_dict)
            self.redis_client.setex(key, record.ttl, record_json)
            
            return True
            
        except Exception as e:
            logger.error(f"保存记录失败: {record.request_id}, error={str(e)}")
            return False
    
    async def cleanup_expired_records(self, batch_size: int = 1000) -> int:
        """清理过期记录"""
        try:
            # 扫描所有幂等性记录
            pattern = f"{self.key_prefix}*"
            keys = self.redis_client.keys(pattern)
            
            expired_count = 0
            current_time = datetime.utcnow()
            
            for key in keys[:batch_size]:  # 限制批处理大小
                try:
                    record_data = self.redis_client.get(key)
                    if record_data:
                        import json
                        record_dict = json.loads(record_data.decode('utf-8'))
                        updated_at = datetime.fromisoformat(record_dict['updated_at'])
                        ttl = record_dict.get('ttl', self.default_ttl)
                        
                        # 检查是否过期
                        if current_time - updated_at > timedelta(seconds=ttl):
                            self.redis_client.delete(key)
                            expired_count += 1
                            
                except Exception as e:
                    logger.error(f"清理记录失败: {key}, error={str(e)}")
                    continue
            
            logger.info(f"清理过期记录完成: 清理了 {expired_count} 条记录")
            return expired_count
            
        except Exception as e:
            logger.error(f"清理过期记录异常: error={str(e)}")
            return 0


# 使用示例
async def example_usage():
    """幂等性管理器使用示例"""
    idempotency_manager = IdempotencyManager()
    
    # 生成request_id
    request_params = {
        'user_id': 'user123',
        'task_type': 'text_qa',
        'query': '什么是人工智能',
        'timestamp': int(time.time())
    }
    
    request_id = idempotency_manager._generate_request_id(**request_params)
    print(f"生成的request_id: {request_id}")
    
    try:
        # 开始处理
        if await idempotency_manager.start_processing(request_id, max_attempts=3):
            print(f"开始处理请求: {request_id}")
            
            # 模拟处理过程
            await asyncio.sleep(1)
            
            # 模拟处理结果
            result = {
                'answer': '人工智能是计算机科学的一个分支...',
                'confidence': 0.95,
                'sources': ['wiki', 'textbook']
            }
            
            # 完成处理
            await idempotency_manager.complete_processing(request_id, result)
            print(f"处理完成: {request_id}")
            
        else:
            print(f"请求已处理过，返回缓存结果")
            
    except IdempotencyError as e:
        print(f"幂等性检查失败: {str(e)}")
        
    except Exception as e:
        # 处理失败
        await idempotency_manager.fail_processing(request_id, str(e))
        print(f"处理失败: {request_id}, error={str(e)}")


if __name__ == "__main__":
    import asyncio
    asyncio.run(example_usage())